import numpy as np
import torch
import torch.nn as nn

class Particle:
    def __init__(self, position, mass, init_time, device):
        self.position = position
        self.mass = mass
        self.time = init_time

        self.positions = []
        self.times = []

        ##
        self.positions.append(self.position)
        self.times.append(self.time)
        self.device = torch.device(device)
        self.velocities = torch.tensor([np.NaN],device=self.device)
        self.accelerations = torch.tensor([np.NaN],device=self.device)
        # self.calculate_velocity()
        # self.calculate_acceleration()


    def add_data(self,position,lateset_time):
        self.positions.append(position)
        self.times.append(lateset_time)

    def interp(self,data_x,data_y,interp_x,left=None,right=None):
        data_new = np.interp(interp_x,data_x,data_y)

        return data_new
        pass

    def calculate_velocity(self):
        # print(np.array(self.positions))
        # print(np.array(self.times))
        self.velocities = torch.diff(torch.tensor(self.positions,device=self.device),dim=0) / torch.diff(torch.tensor(self.times,device=self.device),dim=0)
        # self.velocities = torch.tensor(velocities.clone().detach(),device=self.device)
        self.velocities = torch.vstack([torch.tensor([np.NaN,np.NaN,np.NaN],device=self.device),self.velocities])
        # self.velocities = self.velocities.to(self.device)


    def calculate_acceleration(self):
        # print(np.array(self.positions))
        # print(np.array(self.times))
        self.accelerations = torch.diff(self.velocities,dim=0) / torch.diff(torch.tensor(self.times,device=self.device),dim=0)
        # self.accelerations = torch.tensor(accelerations.clone().detach(),device=self.device)
        self.accelerations = torch.vstack([torch.tensor([np.NaN,np.NaN,np.NaN],device=self.device),self.accelerations])

    def get_attributes(self,time):
        index = self.times.index(time)
        return [self.positions[index],self.velocities[index],self.accelerations[index]]


class Particles:
    def __init__(self, number, positions, masses, velocities=None,init_time=0, mass_scale=1, device="cpu"):
        self.number = number
        self.masses = masses.reshape(self.number,1)
        self.time = init_time

        self.times = []
        self.G = 6.67430e-11 * mass_scale  # 万有引力常数，单位为m^3 kg^-1 s^-2

        ##
        self.positions = positions
        self.times.append(self.time)
        self.device = torch.device(device)
        if velocities is None:
            self.velocities = torch.zeros((self.number,1),device=self.device)
        else:
            self.velocities = velocities

        self.accelerations = torch.zeros((self.number,1),device=self.device)
        # self.calculate_velocity()
        # self.calculate_acceleration()


    def add_time(self,lateset_time):
        # self.positions.append(position)
        self.times.append(lateset_time)

    def interp(self,data_x,data_y,interp_x,left=None,right=None):
        data_new = np.interp(interp_x,data_x,data_y)

        return data_new
        pass

    def update_positions(self):
        self.calculate_acceleration()
        self.calculate_velocity()
        # self.velocities = self.velocities + self.accelerations * (self.times[-1] - self.times[-2])
        if len(self.times) >= 2:
            self.positions = self.positions + self.velocities * (self.times[-1] - self.times[-2])
            if len(self.times) > 2:
                self.times.pop(0)
        self.accelerations = None


    def calculate_velocity(self):
        # print(np.array(self.positions))
        # print(np.array(self.times))
        # self.velocities = torch.diff(torch.tensor(self.positions,device=self.device),dim=0) /\
        #                   torch.diff(torch.tensor(self.times,device=self.device),dim=0)
        # # self.velocities = torch.tensor(velocities.clone().detach(),device=self.device)
        # self.velocities = torch.vstack([torch.tensor([np.NaN,np.NaN,np.NaN],device=self.device),self.velocities])
        # self.velocities = self.velocities.to(self.device)
        self.velocities = self.velocities + self.accelerations * (self.times[-1] - self.times[-2])



    def calculate_acceleration(self):
        # print(np.array(self.positions))
        # print(np.array(self.times))
        # self.accelerations = torch.diff(self.velocities,dim=0) / torch.diff(torch.tensor(self.times,device=self.device),dim=0)
        # # self.accelerations = torch.tensor(accelerations.clone().detach(),device=self.device)
        # self.accelerations = torch.vstack([torch.tensor([np.NaN,np.NaN,np.NaN],device=self.device),self.accelerations])
        self.calculate_gravity()
        # print(self.gravity_sum.shape)
        # print(self.masses.shape)
        self.accelerations = self.gravity_sum / self.masses
        # return

    def calculate_gravity(self):
        # 计算向量的范数
        radius_norm = torch.cdist(self.positions, self.positions, p=2)
        # 计算径向量
        r_vector = self.positions.unsqueeze(0) - self.positions.unsqueeze(1)
        # print(r_vector.shape)
        # print(radius_norm.unsqueeze(-1).shape)
        # 计算单位向量
        radius_unit_vector = r_vector / radius_norm.unsqueeze(-1)

        # print(r_vector)

        # 计算引力的矢量和
        gravity = radius_unit_vector * (self.G * (self.masses @ self.masses.t()) / (radius_norm ** 2)).unsqueeze(-1)

        self.gravity_sum = gravity.nansum(dim=1, keepdim=False)
        # print(self.gravity_sum)
        # return self.gravity.sum(dim=1, keepdim=False)


    def get_attributes(self,time):
        index = self.times.index(time)
        return [self.positions[index],self.velocities[index],self.accelerations[index]]


if __name__ == '__main__':
    particle = Particle([0,0,0],1,0,"cpu")
    particle.add_data([1,1,2],1)
    particle.add_data([1, 1, 2], 2)
    particle.add_data([1, 1, 2], 3)
    particle.calculate_velocity()
    particle.calculate_acceleration()
    # print(particle.velocities)
    # print(particle.accelerations)
    # print(particle.get_attributes(3))
    particles = torch.tensor(
        [[0, 0, 0],
         [2, 2, 1],
         [1, 2, 3],
         [2, 2, 2]
         ],dtype=torch.float32)
    # pdist = nn.PairwiseDistance(p=2)
    # distances = pdist(particles,particles)
    print(particles.shape)
    distances = torch.cdist(particles,particles,p=2)
    print(distances)
    number = 1e3
    device = "cpu"
    position = torch.randn((3,))
    init_time = 0
    mass = 1
    for i in range(int(number)):
        particle = Particle(position,mass,init_time,device=device)
        if i % 1000 == 0:
            print("{}k".format(i+1))